Add code for AbstractMCMC Callbacks support - #2755

Turing.jl documentation for PR #2755 is available at:
https://TuringLang.github.io/Turing.jl/previews/PR2755/

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 85.17%. Comparing base (e487ed0) to head (58a4496).
⚠️ Report is 2 commits behind head on main.

@@            Coverage Diff             @@
##             main    #2755      +/-   ##
==========================================
+ Coverage   85.12%   85.17%   +0.05%     
==========================================
  Files          20       20              
  Lines        1304     1309       +5     
==========================================
+ Hits         1110     1115       +5     
  Misses        194      194              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

Thanks @shravanngoswamii — I’ve added a few refactoring comments above. The functionality is already in place, but it could benefit from some minor cleanup, which I agree with Penny on.

CI failure is not relevent to this PR.

Thanks @shravanngoswamii -- this is looking nice and clean now!

There is one point that I think we should discuss, which is the meaning of the 'params' and 'stats'. In particular, for samplers which operate in linked space, what do we return? e.g. does getparams return a vector where x is always between 0 or 1? (if this is confusing, there's an explanation here https://turinglang.org/docs/developers/transforms/distributions/)

@model function f()
    x ~ Beta(2, 2)
end

Currently, the implementation of getparams and getstats returns whatever the sampler stores internally. That means that some samplers will have unlinked values (e.g. default MH, which samples statically from the prior in unlinked space) and some samplers will have linked values (e.g. NUTS). The same is true of stats -- currently we return the sampler's notion of lp which may or may not be linked.

I don't mind as long as we are consistent, for example, we could say that getstats and getparams for Turing samplers always return whatever is internal to the sampler -- but I do think that it is much better from a user perspective to get the unlinked values since the internal values have no external meaning. We used to return internal log probs in the MCMCChains object, but we switched it over to returning 'true' unlinked lps recently, and I think that was a good thing.

Thank you @penelopeysm.

getparams(state) returns get_varinfo(state)[:]:

• HMC/NUTS/SGHMC/SGLD: Linked (unconstrained) -- these samplers explicitly link the VarInfo
• MH() (default/static proposals): Unlinked -- should_link returns false
• MH(covariance_matrix) or explicit RandomWalkProposal: Linked
• ESS/PG: Unlinked

getstats(state).lp is also inconsistent:

• HMC: sum(getlogp(vi)) = logprior + logjac + loglikelihood
• MH: logjoint_internal = logprior + loglikelihood - logjac

However, ParamsWithStats (the transition) already does the right thing -- DynamicPPL.ParamsWithStats re-evaluates the model via ValuesAsInModelAccumulator and returns unlinked parameter values with proper logjoint (no Jacobian artifacts).

For callback-based monitoring, users interact with ParamsWithStats, not the raw state. So the user-facing behavior is correct.

Maybe, we can document that getparams/getstats return internal sampler state (which may be in transformed space), while ParamsWithStats provides semantically correct values. This matches the intent that these are low-level accessors vs. high-level transition containers.

Happy to add this documentation if that sounds reasonable.

Tests will pass after TuringLang/AbstractMCMC.jl#191

> julia --project test\mcmc\callbacks.jl          
┌ Info: Found initial step size
└   ϵ = 1.6500000000000001
┌ Info: Found initial step size
└   ϵ = 1.6500000000000001
Test Summary:                    | Pass  Total   Time
AbstractMCMC Callbacks Interface |   80     80  38.3s

However, ParamsWithStats (the transition) already does the right thing -- DynamicPPL.ParamsWithStats re-evaluates the model via ValuesAsInModelAccumulator and returns unlinked parameter values with proper logjoint (no Jacobian artifacts). For callback-based monitoring, users interact with ParamsWithStats, not the raw state. So the user-facing behavior is correct.

Does it? Don't all the methods getstats(...), etc., act on the state, not the transition? And that's what the callback framework uses?

Does it? Don't all the methods getstats(...), etc., act on the state, not the transition? And that's what the callback framework uses?

I think the callback behavior is correct. The TensorBoard callback in AbstractMCMC calls ParamsWithStats(model, sampler, transition, state) and Turing overrides this at Inference.jl for transition::DynamicPPL.ParamsWithStats.

This override extracts transition.params and transition.stats directly—it doesn't call getparams(state) or getstats(state).

Right now, the ParamsWithStats override:

• Extracts params from transition.params
• Extracts stats from transition.stats (includes logprior, loglikelihood, logjoint, and AHMC diagnostics like tree_depth, n_steps, acceptance_rate)
• Returns empty extras (always NamedTuple())
• Ignores state entirely

I am not sure if ignoring state is acceptable? Is the current design acceptable, or should we populate extras with static sampler configuration from state?
What are your thoughts @penelopeysm @yebai ?

Since it's nonbreaking to add extra info in, I would be OK with this current implementation right now, and afterwards if there's a need for more information that can be changed later / in a separate PR.

I’m happy with the PR!

Why the minor version bump?

Why the minor version bump?

New funtion overloaded should be patch or minor?

It's backwards compatible new functionality. If the version number is 0.x, it only needs to be a patch bump. If the version number is N.x (N ≥ 1) it only needs to be a minor bump.

But semver says you do have a choice: if you prefer to release it as a minor version 0.(x+1) or a major version (N+1), you can.

If you do choose to release as a minor version, please can you rebase it against breaking? See description in https://turinglang.org/docs/developers/contributing/#pull-requests-versions-and-releases (although we only really follow this for Turing and DPPL right now)

I would recommend not doing it as a minor version though, as the next minor version will take a while to be done, since I have to do #2756 and that is very difficultl

But semver says you do have a choice: if you prefer to release it as a minor version 0.(x+1) or a major version (N+1), you can.

Okay, I get it! Thanks for all the reviews, @penelopeysm. Let me know if you’re okay with merging this.